# @Author : wangyuchen
# @Time : 2021-05-15 14:01

import argparse


import torch

from utils import settings
from utils.base import get_network, get_testing_dataloader

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-net', type=str, required=True, help='net type')
    parser.add_argument('-weights', type=str, required=True, help='the weights file you want to test')
    parser.add_argument('-gpu', type=bool, default=False, help='use gpu or not')
    parser.add_argument('-b', type=int, default=16, help='batch size for dataloader')
    args = parser.parse_args()
    net = get_network(args)
    testing_loader = get_testing_dataloader(
        './test',
        settings.VAL_MEAN,
        settings.VAL_STD,
        num_workers=4,
        batch_size=args.b,
        shuffle=True
    )
    if args.gpu:
        net.load_state_dict(torch.load(args.weights))
    else:
        net.load_state_dict(torch.load(args.weights, map_location='cpu'))
    net.eval()
    fp = open(settings.TEST_RES_PATH.format(net=args.net), 'w+')
    with torch.no_grad():
        for n_iter, (images, img_names) in enumerate(testing_loader):
            print("iteration: {}\ttotal {} iterations".format(n_iter + 1, len(testing_loader)))
            if args.gpu:
                images = images.cuda()
            outputs = net(images)
            _, preds = torch.max(outputs.data, 1)
            for i in range(len(preds)):
                fp.write(img_names[i] + ' ' + str(preds[i].item()) + '\n')
    fp.close()
